package com.core.cloudcommon.trace;

import lombok.extern.slf4j.Slf4j;

import java.util.UUID;

/**
 * @author =====Lqj=====
 * @className TraceContext
 * @description 链路追踪上下文管理器
 * @create 2024/12/17 10:30
 **/
@Slf4j
public class TraceContext {

    /**
     * 链路ID的ThreadLocal存储
     */
    private static final ThreadLocal<String> TRACE_ID_HOLDER = new ThreadLocal<>();

    /**
     * 用户ID的ThreadLocal存储
     */
    private static final ThreadLocal<Long> USER_ID_HOLDER = new ThreadLocal<>();

    /**
     * 用户名的ThreadLocal存储
     */
    private static final ThreadLocal<String> USER_NAME_HOLDER = new ThreadLocal<>();

    /**
     * 请求URI的ThreadLocal存储
     */
    private static final ThreadLocal<String> REQUEST_URI_HOLDER = new ThreadLocal<>();

    /**
     * 客户端IP的ThreadLocal存储
     */
    private static final ThreadLocal<String> CLIENT_IP_HOLDER = new ThreadLocal<>();

    /**
     * 请求开始时间的ThreadLocal存储
     */
    private static final ThreadLocal<Long> START_TIME_HOLDER = new ThreadLocal<>();

    /**
     * 链路ID的HTTP头名称
     */
    public static final String TRACE_ID_HEADER = "X-Trace-Id";

    /**
     * 生成新的链路ID
     * @return 链路ID
     */
    public static String generateTraceId() {
        return UUID.randomUUID().toString().replace("-", "");
    }

    /**
     * 设置链路ID
     * @param traceId 链路ID
     */
    public static void setTraceId(String traceId) {
        TRACE_ID_HOLDER.set(traceId);
        // 同步到MDC，让所有日志自动包含链路ID
        TraceMDCManager.setTraceId(traceId);
        log.debug("设置链路ID: {}", traceId);
    }

    /**
     * 获取当前链路ID
     * @return 链路ID
     */
    public static String getTraceId() {
        return TRACE_ID_HOLDER.get();
    }

    /**
     * 设置用户ID
     * @param userId 用户ID
     */
    public static void setUserId(Long userId) {
        USER_ID_HOLDER.set(userId);
        // 同步到MDC
        TraceMDCManager.setUserId(userId);
    }

    /**
     * 获取当前用户ID
     * @return 用户ID
     */
    public static Long getUserId() {
        return USER_ID_HOLDER.get();
    }

    /**
     * 设置用户名
     * @param userName 用户名
     */
    public static void setUserName(String userName) {
        USER_NAME_HOLDER.set(userName);
        // 同步到MDC
        TraceMDCManager.setUserName(userName);
    }

    /**
     * 获取当前用户名
     * @return 用户名
     */
    public static String getUserName() {
        return USER_NAME_HOLDER.get();
    }

    /**
     * 设置请求URI
     * @param requestUri 请求URI
     */
    public static void setRequestUri(String requestUri) {
        REQUEST_URI_HOLDER.set(requestUri);
        // 同步到MDC
        TraceMDCManager.setRequestUri(requestUri);
    }

    /**
     * 获取当前请求URI
     * @return 请求URI
     */
    public static String getRequestUri() {
        return REQUEST_URI_HOLDER.get();
    }

    /**
     * 设置客户端IP
     * @param clientIp 客户端IP
     */
    public static void setClientIp(String clientIp) {
        CLIENT_IP_HOLDER.set(clientIp);
        // 同步到MDC
        TraceMDCManager.setClientIp(clientIp);
    }

    /**
     * 获取当前客户端IP
     * @return 客户端IP
     */
    public static String getClientIp() {
        return CLIENT_IP_HOLDER.get();
    }

    /**
     * 设置请求开始时间
     * @param startTime 开始时间
     */
    public static void setStartTime(Long startTime) {
        START_TIME_HOLDER.set(startTime);
    }

    /**
     * 获取请求开始时间
     * @return 开始时间
     */
    public static Long getStartTime() {
        return START_TIME_HOLDER.get();
    }

    /**
     * 计算请求耗时
     * @return 耗时（毫秒）
     */
    public static Long getElapsedTime() {
        Long startTime = getStartTime();
        if (startTime != null) {
            return System.currentTimeMillis() - startTime;
        }
        return null;
    }

    /**
     * 初始化链路追踪上下文
     * @param traceId 链路ID（可选，为空则自动生成）
     * @param userId 用户ID
     * @param userName 用户名
     * @param requestUri 请求URI
     * @param clientIp 客户端IP
     */
    public static void initTrace(String traceId, Long userId, String userName, 
                                String requestUri, String clientIp) {
        // 如果没有传入链路ID，则生成新的
        if (traceId == null || traceId.trim().isEmpty()) {
            traceId = generateTraceId();
        }
        
        setTraceId(traceId);
        setUserId(userId);
        setUserName(userName);
        setRequestUri(requestUri);
        setClientIp(clientIp);
        setStartTime(System.currentTimeMillis());
        
//        log.info("初始化链路追踪 - TraceId: {}, UserId: {}, UserName: {}, URI: {}, IP: {}",
//                traceId, userId, userName, requestUri, clientIp);
    }

    /**
     * 清理链路追踪上下文
     */
    public static void clear() {
        String traceId = getTraceId();
        Long elapsedTime = getElapsedTime();

        TRACE_ID_HOLDER.remove();
        USER_ID_HOLDER.remove();
        USER_NAME_HOLDER.remove();
        REQUEST_URI_HOLDER.remove();
        CLIENT_IP_HOLDER.remove();
        START_TIME_HOLDER.remove();

        // 清理MDC
        TraceMDCManager.clear();

//        log.info("清理链路追踪 - TraceId: {}, 耗时: {}ms", traceId, elapsedTime);
    }

    /**
     * 获取当前链路追踪信息
     * @return 链路追踪信息
     */
    public static TraceInfo getCurrentTraceInfo() {
        return TraceInfo.builder()
                .traceId(getTraceId())
                .userId(getUserId())
                .userName(getUserName())
                .requestUri(getRequestUri())
                .clientIp(getClientIp())
                .startTime(getStartTime())
                .elapsedTime(getElapsedTime())
                .build();
    }

    /**
     * 检查是否存在链路追踪上下文
     * @return true-存在，false-不存在
     */
    public static boolean hasTrace() {
        return getTraceId() != null;
    }

    /**
     * 链路追踪信息实体
     */
    public static class TraceInfo {
        private String traceId;
        private Long userId;
        private String userName;
        private String requestUri;
        private String clientIp;
        private Long startTime;
        private Long elapsedTime;

        public static TraceInfoBuilder builder() {
            return new TraceInfoBuilder();
        }

        // Getters
        public String getTraceId() { return traceId; }
        public Long getUserId() { return userId; }
        public String getUserName() { return userName; }
        public String getRequestUri() { return requestUri; }
        public String getClientIp() { return clientIp; }
        public Long getStartTime() { return startTime; }
        public Long getElapsedTime() { return elapsedTime; }

        @Override
        public String toString() {
            return String.format("TraceInfo{traceId='%s', userId=%d, userName='%s', requestUri='%s', clientIp='%s', startTime=%d, elapsedTime=%d}", 
                    traceId, userId, userName, requestUri, clientIp, startTime, elapsedTime);
        }

        public static class TraceInfoBuilder {
            private String traceId;
            private Long userId;
            private String userName;
            private String requestUri;
            private String clientIp;
            private Long startTime;
            private Long elapsedTime;

            public TraceInfoBuilder traceId(String traceId) {
                this.traceId = traceId;
                return this;
            }

            public TraceInfoBuilder userId(Long userId) {
                this.userId = userId;
                return this;
            }

            public TraceInfoBuilder userName(String userName) {
                this.userName = userName;
                return this;
            }

            public TraceInfoBuilder requestUri(String requestUri) {
                this.requestUri = requestUri;
                return this;
            }

            public TraceInfoBuilder clientIp(String clientIp) {
                this.clientIp = clientIp;
                return this;
            }

            public TraceInfoBuilder startTime(Long startTime) {
                this.startTime = startTime;
                return this;
            }

            public TraceInfoBuilder elapsedTime(Long elapsedTime) {
                this.elapsedTime = elapsedTime;
                return this;
            }

            public TraceInfo build() {
                TraceInfo traceInfo = new TraceInfo();
                traceInfo.traceId = this.traceId;
                traceInfo.userId = this.userId;
                traceInfo.userName = this.userName;
                traceInfo.requestUri = this.requestUri;
                traceInfo.clientIp = this.clientIp;
                traceInfo.startTime = this.startTime;
                traceInfo.elapsedTime = this.elapsedTime;
                return traceInfo;
            }
        }
    }
}
